# -*- coding: utf-8 -*-
"""
Created on Thu Jun 10 21:28:40 2021

@author: Administrator
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
from sklearn.linear_model import LinearRegression
from scipy.stats import pearsonr

plt.rcParams['font.sans-serif']=['SimHei']    #画图时使用中文字体
plt.rcParams['axes.unicode_minus'] = False

def lag_plotting(series, lag=1, residual=False):
    '''
    画出 observation and previous observation with specific lag
    '''
    fig = plt.figure()
    pd.plotting.lag_plot(series, lag=lag)
    plt.xlabel('数据（t）', fontsize=16)
    plt.ylabel(f'数据(t-{lag})', fontsize=16)
    plt.title('当前数据和延迟数据', fontsize=24)
    if residual:
        # residual == True 时，绘制残差图
        fig.savefig(f'../图片/残差数据与时间延迟{lag}的数据.png')
    else:    
        fig.savefig(f'../图片/当前数据与时间延迟{lag}的数据.png')
    plt.grid()
    plt.show()
    

def calculate_corr(series, lag=5):
    '''
    calculate the autocorrlation of observations and previous data with specific lags
    '''
    values = pd.DataFrame(series.values)
    data = [values.shift(i) for i in range(0,lag+1)]
    columns = [f'x({i})' for i in range(0,lag+1)]
    indexs = columns
    dataframe = pd.concat(data, axis=1)
    # calculate autocorrlation matrix
    corr = dataframe.corr()
    corr.columns = columns
    corr.index = indexs
    return corr
    

def calculate_partial_corr(series, lag=5):
    '''
    calculate paritial autocorrelation matrix
    '''
    values = pd.DataFrame(series.values)
    data = [values.shift(i) for i in range(0,lag+1)]
    columns = [f'x({i})' for i in range(0,lag+1)]
    df = pd.concat(data, axis=1)
    # 删除包含 nan 的行
    df.dropna(inplace=True)
    df.columns = columns
    # 初始化偏相关矩阵：
    partial_corr_matrix = np.zeros((lag+1, lag+1))
    for i in range(lag+1):
        x1 = df.iloc[:, i]
        for j in range(lag+1):
            if i == j:
                partial_corr_matrix[i, j] = 1
            elif j < i:
                partial_corr_matrix[i, j] = partial_corr_matrix[j, i]
            else:
                x2 = df.iloc[:, j]
                df_control = df.drop(columns=[columns[i], columns[j]], axis=1, inplace=False)
                L = LinearRegression().fit(df_control, x1)
                Lx = L.predict(df_control)
                x1_prime = x1 - Lx
                
                L = LinearRegression().fit(df_control, x2)
                Lx = L.predict(df_control)
                x2_prime = x2 - Lx
                partial_corr_matrix[i, j] = pearsonr(x1_prime, x2_prime)[0]
    
    indexs = columns
    partial_corr = pd.DataFrame(partial_corr_matrix, columns=columns)
    partial_corr.index = indexs
    return partial_corr

def acf_plotting(series, lags=30, residual=False):
    fig = plt.figure()
    plot_acf(series, lags=lags, alpha=0.05)
    plt.xlabel('时间延迟 t', fontsize=16)
    plt.ylabel('自相关系数', fontsize=16)
    plt.title('自相关图', fontsize=24)
    if residual:
        fig.savefig(f'../图片/残差数据的自相关图.png')
    else:
        fig.savefig(f'../图片/自相关图.png')
    plt.grid()
    plt.show()

def pacf_plotting(series, lags=30, residual=False):
    fig = plt.figure()
    plot_pacf(series, lags=lags, alpha=0.05)
    plt.xlabel('时间延迟 t', fontsize=16)
    plt.ylabel('偏相关系数', fontsize=16)
    plt.title('偏相关图', fontsize=24)
    if residual:
        fig.savefig(f'../图片/残差数据的自相关图.png')
    else:
        fig.savefig(f'../图片/偏相关图.png')
    plt.grid()
    plt.show()
    

def analyze_matrix_drawing(series):
    lag_plotting(series, lag=1)
    lag_plotting(series, lag=5)
    corr = calculate_corr(series, lag=5)
    print('自相关系数矩阵：\n', corr)
    partial_corr = calculate_partial_corr(series, lag=5)
    print('偏相关系数矩阵：\n', partial_corr)
    acf_plotting(series, lags=30)
    pacf_plotting(series, lags=30)
   
if __name__ == '__main__':
    path = r'../附件/中间数据/光伏-建筑板块市值.xlsx'
    series = pd.read_excel(path, header=0, index_col=0, parse_dates=True, squeeze=True)
    series = series['2019-4-1':'2021-4-30']
    acf_plotting(series)
    pacf_plotting(series)
    corr = calculate_corr(series)
    # partial_corr = calculate_partial_corr(series)
    lag_plotting(series, lag=1)
    lag_plotting(series, lag=7)
    lag_plotting(series, lag=30)
    print(corr)
    # print(partial_corr)